# Copyright 2024 DeepMind Technologies Limited
#
# AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
# this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
#
# To request access to the AlphaFold 3 model parameters, follow the process set
# out at https://github.com/google-deepmind/alphafold3. You may only use these
# if received directly from Google. Use is subject to terms of use available at
# https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md

"""Library to run Mmseqs from Python."""

import os
import tempfile
import time
import shutil

from absl import logging
from alphafold3.data import parsers
from alphafold3.data.tools import msa_tool
from alphafold3.data.tools import subprocess_utils
import shlex

#  --msa-format-mode INT   Format MSA as: 0: binary cA3M DB
#                          1: binary ca3m w. consensus DB
#                          2: aligned FASTA DB
#                          3: aligned FASTA w. header summary
#                          4: STOCKHOLM flat file
#                          5: A3M format
#                          6: A3M format w. alignment info [2]

class Mmseqs(msa_tool.MsaTool):
  """Python wrapper of the Mmseqs binary."""

  def __init__(
      self,
      *,
      binary_path: str,
      database_path: str,
      use_gpu: int=0,
      msa_format_mode: int=4,
      e_value: float | None = 1e-3,
      max_sequences: int = 5000,
      mmseqs_options: str,
      result2msa_options: str,
      gpu_id: int=0,
      threads: int=8
  ):
    """Initializes the Python Mmseqs wrapper.

    Args:

    """
    self.binary_path = binary_path
    self.database_path = database_path
    self.mmseqs_options = mmseqs_options
    self.result2msa_options = result2msa_options
    self.gpu_id = gpu_id

    subprocess_utils.check_binary_exists(
        path=self.binary_path, name='Mmseqs'
    )

    if not os.path.exists(self.database_path):
      raise ValueError(f'Could not find Mmseqs database {database_path}')

    self.use_gpu = use_gpu
    self.msa_format_mode = msa_format_mode
    self.e_value = e_value
    self.max_sequences = max_sequences
    self.threads = threads

  def query(self, target_sequence: str) -> msa_tool.MsaToolResult:
    """Queries the database using MMseqs."""
    logging.info('Query sequence: %s', target_sequence)
    os.makedirs('/tmp_owm', exist_ok=True)
    with tempfile.TemporaryDirectory(dir='/tmp_owm') as query_tmp_dir:
      input_fasta_path = os.path.join(query_tmp_dir, 'query.fasta')
      subprocess_utils.create_query_fasta_file(
          sequence=target_sequence, path=input_fasta_path
      )

      logging.info('Searching %s', self.database_path)
      search_start_time = time.time()
      
      sequence_db_path = os.path.join(query_tmp_dir, 'queryDB')
      self._run_createdb_command(input_fasta_path, sequence_db_path)

      result_db_path = os.path.join(query_tmp_dir, 'resultDB')
      self._run_search_command(sequence_db_path, result_db_path)

      output_sto_path = os.path.join(query_tmp_dir, 'output.sto')
      self._run_result2msa_command(sequence_db_path, result_db_path, output_sto_path)

    #   destination_dir = os.path.abspath("./")  # 当前目录绝对路径
    #   new_filename = f"{os.path.splitext(os.path.basename(self.database_path))[0]}.sto"  # 新文件名
    #   dest_path = os.path.join(destination_dir, new_filename)  # 组合完整目标路径
    #   shutil.copy(output_sto_path, dest_path) 

      with open(output_sto_path) as f:
        try:
            # 尝试使用convert_stockholm_to_a3m函数
            a3m = parsers.convert_stockholm_to_a3m(
                f,
                max_sequences=self.max_sequences
            )
        except Exception as e:  # 捕获convert_stockholm_to_a3m可能抛出的任何异常
            print(f"convert_stockholm_to_a3m failed with error: {e}. Trying convert_mmseqs_stockholm_to_a3m now.")
            
            # 重置文件指针
            f.seek(0)

            try:
                # 尝试使用convert_mmseqs_stockholm_to_a3m函数
                a3m = parsers.convert_mmseqs_stockholm_to_a3m(
                    f,
                    max_sequences=self.max_sequences
                )
            except Exception as e:  # 捕获convert_mmseqs_stockholm_to_a3m可能抛出的任何异常
                print(f"convert_mmseqs_stockholm_to_a3m failed with error: {e}. Trying convert_mmseqs_stockholm_to_a3m_unique now.")
            
                # 重置文件指针
                f.seek(0)
            
                # convert_mmseqs_stockholm_to_a3m_unique
                a3m = parsers.convert_mmseqs_stockholm_to_a3m_unique(
                    f,
                    max_sequences=self.max_sequences
                )

    logging.info(
    'Searching took %.2f seconds from %s',
    time.time() - search_start_time,
    self.database_path,
    )
    # with open(f"{os.path.splitext(os.path.basename(self.database_path))[0]}.a3m", "w") as f:
    #     f.write(a3m)

    return msa_tool.MsaToolResult(
        target_sequence=target_sequence, a3m=a3m, e_value=self.e_value
    )
    
  def _run_createdb_command(self, input_fasta_path: str, output_db_path: str):
        """Runs the MMseqs2 `createdb` command."""
        cmd = [
            self.binary_path,
            'createdb',
            input_fasta_path,
            output_db_path,
        ]
        subprocess_utils.run(
            cmd=cmd,
            cmd_name='MMseqs2 createdb',
            log_stdout=True,
            log_stderr=True,
            log_on_process_error=True,
        )

  def _run_search_command(self, input_db_path: str, output_db_path: str):
        """Runs the MMseqs2 `search` command."""

#         --split INT                      将输入分割成 N 个均匀分布的块。0: 自动设置最佳分割 [0]
#  --split-mode INT                 0: 分割目标数据库；1: 分割查询数据库；2: 自动，取决于主内存 [2]
#  --split-memory-limit BYTE        设置每个分割的最大内存。例：800B, 5K, 10M, 1G。默认（0）为所有可用系统内存 [0]
# --threads INT                    使用的 CPU 核心数量（默认全部） [32]
      # --prefilter-mode INT             预过滤模式：0: kmer/无缺口 1: 无缺口, 2: 不过滤, 3: 无缺口&有缺口 [0]
        
        env = os.environ.copy()
        # if self.use_gpu:
        #     # env["HIP_VISIBLE_DEVICES"] = str(self.gpu_id)
        #     self.threads = 32

        cmd = [
            self.binary_path,
            'search',
            input_db_path,
            self.database_path,
            output_db_path,
            '/tmp_owm',  # Temporary directory for intermediate files
            '--gpu',
            str(self.use_gpu),
            '--threads',str(self.threads),
        ]
        cmd.extend(shlex.split(self.mmseqs_options))
        subprocess_utils.run(
            cmd=cmd,
            env=env,
            cmd_name='MMseqs2 search',
            log_stdout=True,
            log_stderr=True,
            log_on_process_error=True,
        )

  def _run_result2msa_command(self, input_db_path: str, result_db_path: str, output_msa_path: str):
        """Runs the MMseqs2 `result2msa` command."""
        cmd = [
            self.binary_path,
            'result2msa',
            input_db_path,
            self.database_path,
            result_db_path,
            output_msa_path,
            '--msa-format-mode',
            str(self.msa_format_mode),
        ]
        cmd.extend(shlex.split(self.result2msa_options))
        subprocess_utils.run(
            cmd=cmd,
            cmd_name='MMseqs2 result2msa',
            log_stdout=True,
            log_stderr=True,
            log_on_process_error=True,
        )

